import os
import logging
from scapy.layers.inet import IP, TCP, UDP
from scapy.layers.l2 import Ether
from scapy.all import sniff, sendp, Raw, wrpcap
from threading import Thread, Event

# 定义全局变量
capture_running = Event()
packets = []
log_path = "xieyi"  # 默认日志路径
save_path = "xieyi"  # 默认保存数据包路径
logger = logging.getLogger()

# 配置日志
def configure_logging(log_path):
    log_file = os.path.join(log_path, "network_tool.log")
    logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    logger.info(f"日志路径已设置为: {log_path}")

configure_logging(log_path)

def packet_callback(packet):
    if packet.haslayer(IP):
        ip_src = packet[IP].src
        ip_dst = packet[IP].dst
        if packet.haslayer(TCP):
            tcp_sport = packet[TCP].sport
            tcp_dport = packet[TCP].dport
            logger.info(f"TCP Packet: {ip_src}:{tcp_sport} -> {ip_dst}:{tcp_dport}")
            print(f"TCP Packet: {ip_src}:{tcp_sport} -> {ip_dst}:{tcp_dport}")
        elif packet.haslayer(UDP):
            udp_sport = packet[UDP].sport
            udp_dport = packet[UDP].dport
            logger.info(f"UDP Packet: {ip_src}:{udp_sport} -> {ip_dst}:{udp_dport}")
            print(f"UDP Packet: {ip_src}:{udp_sport} -> {ip_dst}:{udp_dport}")
        packets.append(packet)

def start_capture(sources, filter_rule, log_path, packets_list):
    global capture_running
    capture_running.set()
    ifaces = []
    if "1" in sources:  # 数传服务器
        ifaces.append("eth1")
    if "2" in sources:  # 本地
        ifaces.append("lo")
    if "3" in sources:  # 网络
        ifaces.append("eth0")

    def capture_thread(iface):
        sniff(iface=iface, prn=packet_callback, filter=filter_rule, stop_filter=lambda x: not capture_running.is_set())

    threads = []
    for iface in ifaces:
        thread = Thread(target=capture_thread, args=(iface,))
        thread.daemon = True  # 设置为守护线程
        thread.start()
        threads.append(thread)

    # 等待所有线程结束
    for thread in threads:
        thread.join()

    # 将抓取到的包传递给前端
    packets_list.extend(packets)

def stop_capture():
    global capture_running
    capture_running.clear()

def send_packet(destinations, protocol, data, log_path):
    src_ip = "127.0.0.1"  # 默认本地发送
    for destination in destinations:
        if destination == "1":  # 数传服务器
            dst_ip = "192.168.1.201"
        elif destination == "2":  # 本地
            dst_ip = "127.0.0.1"
        elif destination == "3":  # 网络
            dst_ip = "192.168.1.202"

        if protocol == "TCP":
            packet = Ether() / IP(src=src_ip, dst=dst_ip) / TCP() / Raw(load=data)
        elif protocol == "UDP":
            packet = Ether() / IP(src=src_ip, dst=dst_ip) / UDP() / Raw(load=data)

        try:
            sendp(packet)
            logger.info(f"成功发送 {protocol} 包: 本地 -> {destination}")
            print(f"成功发送 {protocol} 包: 本地 -> {destination}")
        except Exception as e:
            logger.error(f"发送 {protocol} 包失败: 本地 -> {destination}, 错误信息: {e}")
            print(f"发送 {protocol} 包失败: 本地 -> {destination}, 错误信息: {e}")

def save_packets(save_path, packets_list):
    save_file = os.path.join(save_path, "captured_packets.pcap")
    wrpcap(save_file, packets_list)
    logger.info(f"抓取到的包已保存到文件: {save_file}")
    print(f"抓取到的包已保存到文件: {save_file}")

def set_log_path(new_log_path):
    global log_path
    log_path = os.path.abspath(new_log_path)
    configure_logging(log_path)
    logger.info(f"日志路径已更新为: {log_path}")